from typing import List


class Solution:
    def removeBoxes(self, boxes: List[int]) -> int:
        # 处理只有1个盒子的特殊情况
        if len(boxes) == 1:
            return 1

        # 压缩盒子颜色
        index = {box: i for i, box in enumerate(sorted(set(boxes)))}
        boxes = [index[box] for box in boxes]

        # 压缩盒子列表
        # groups = []
        # for box in boxes:
        #     if not groups or groups[-1][0] != box:
        #         groups.append([box, 1])
        #     else:
        #         groups[-1][1] += 1
        #
        # print("盒子组合列表:", groups)

        box_num = len(boxes)  # 盒子数量
        # group_num = len(groups)  # 盒子组合数量
        color_num = len(index)  # 盒子颜色数量

        print("盒子列表:", boxes, ";", "盒子颜色数量:", color_num)

        # 定义状态矩阵：dp[i][j][k]=取走前i个盒子，最后取走的颜色为j、最后一次取走了k个的最大积分和
        dp = [[[0 for _ in range(box_num)] for _ in range(color_num)] for _ in range(box_num)]
        dp[0][boxes[0]][1] = 1

        for color, row in enumerate(dp[0]):
            print(color, ":", row)

        # 状态转移
        for i in range(1, box_num):
            color1 = boxes[i]  # 新盒子颜色
            for color2 in range(color_num):  # 遍历最后取走的旧盒子的颜色
                if color1 == color2:  # 最后取走的旧盒子与新盒子颜色相同
                    for j in range(box_num - 1, 0, -1):
                        if dp[i - 1][color2][j - 1]:
                            dp[i][color1][j] = max(dp[i][color1][j], dp[i - 1][color2][j - 1] - (j - 1) ** 2 + j ** 2)
                else:  # 最后取走的旧盒子与新盒子颜色不同
                    for j in range(box_num - 1, 0, -1):
                        if dp[i - 1][color2][j]:
                            dp[i][color1][1] = max(dp[i][color1][1], dp[i - 1][color2][j] + 1)  # 先取走旧盒子
                            dp[i][color2][j] = max(dp[i][color2][j], dp[i - 1][color2][j] + 1)  # 先取走新盒子

            print("当前盒子:", i, "颜色:", color1)
            for color, row in enumerate(dp[i]):
                print(color, ":", row)

        # 计算结果
        ans = 0
        for color in range(color_num):
            for j in range(1, box_num):
                ans = max(ans, dp[-1][color][j])
        return ans


if __name__ == "__main__":
    # [0, 2, 1, 1, 1, 2, 3, 2, 0]
    print(Solution().removeBoxes([1, 3, 2, 2, 2, 3, 4, 3, 1]))  # 23

    # print(Solution().removeBoxes([1, 1, 1]))  # 9
    # print(Solution().removeBoxes([1]))  # 1

    # 自制用例
    # print(Solution().removeBoxes([1, 2, 3, 4, 5, 6, 1, 6, 5, 4, 3, 2, 1]))  # 25
